/*
   Copyright 2009 Abdulla Abdurakhmanov (abdulla@latestbit.com)
   Original sources are available at www.latestbit.com

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

#include "cpputils/CppUtils.hpp"
#include "cpputils/ept/ThreadPool.hpp"

namespace cpputils {
	namespace ept {

		ThreadPool::ThreadPool ( unsigned int nInitialThreadCount, unsigned int nMaxThreadCount, const std::string& sThreadPoolName ) :
					m_nInitialThreadCount (nInitialThreadCount),
					m_nMaxThreadCount(nMaxThreadCount),
					m_nCurrentThreadCount(0),
					m_nTaskStorageSize (0),
					m_nTaskIdCounter (1),
					m_bStopFlag ( false ),
					m_sThreadPoolName(sThreadPoolName) {

			if(nInitialThreadCount > nMaxThreadCount)
				throw std::invalid_argument("Invalid initial thread count");
			if(nMaxThreadCount == 0 )
				throw std::invalid_argument("Invalid max thread count");
			m_poThreadGroup = new boost::thread_group();
			createInitialThreads();
		}

		ThreadPool::~ThreadPool() {
			boost::recursive_mutex::scoped_lock oLock (m_mxTaskStorage);
			m_bStopFlag = true;
			m_evTaskAvailables.notify_all();
			oLock.unlock();

			m_poThreadGroup->join_all();
			delete m_poThreadGroup;
			destroySystemTaskIds();
			destroyMonitors();
		}

		void ThreadPool::createInitialThreads(void) {
			for(unsigned int i=0;i<m_nInitialThreadCount;i++) {
				m_poThreadGroup->create_thread(
					boost::bind(&ThreadPool::onRunTask, this, createSystemTaskId())
				);
				m_nCurrentThreadCount++;
			}
		}

		void ThreadPool::destroyMonitors() {
			boost::recursive_mutex::scoped_lock oLock (m_mxTaskMonitors);
			for(TaskMonitors::iterator it = m_mTaskMonitors.begin(); it!= m_mTaskMonitors.end(); it++) {
				it->second.destroy();
			}
			m_mTaskMonitors.clear();
		}

		ThreadPool::TaskId ThreadPool::createTaskId() {
			boost::recursive_mutex::scoped_lock oLock (m_mxTaskMonitors);
			if(!m_nTaskIdCounter)
				m_nTaskIdCounter = 1; // reserved task id = 0

			// Creating monitor for task id
			std::pair<TaskId, TaskMonitor> oPair;
			oPair.first = m_nTaskIdCounter;
			oPair.second =  TaskMonitor( this );
			m_mTaskMonitors.insert(oPair);

			return m_nTaskIdCounter++;
		}

		ThreadPool::TaskId ThreadPool::createSystemTaskId() {
			boost::recursive_mutex::scoped_lock oLock (m_mxSystemTaskId);
			TaskId nSystemTaskId = createTaskId();
			m_itCurrentTaskId = m_aSystemTaskIds.insert ( m_aSystemTaskIds.end(), nSystemTaskId );
			return nSystemTaskId;
		}

		void ThreadPool::destroySystemTaskIds() {
			boost::recursive_mutex::scoped_lock oLock (m_mxSystemTaskId);
			foreach (TaskId nTaskId, m_aSystemTaskIds) {
				destroyTaskId ( nTaskId );
			}
			m_aSystemTaskIds.clear();
		}

		void ThreadPool::destroyTaskId(TaskId nTaskId) {
			boost::recursive_mutex::scoped_lock oLock (m_mxTaskMonitors);
			cancel(nTaskId);
			wait(nTaskId);

			// Removing monitor for task id
			TaskMonitors::iterator it = m_mTaskMonitors.find(nTaskId);
			if(it!=m_mTaskMonitors.end()) {
				boost::recursive_mutex::scoped_lock oTaskMonitorLock ( it->second.getMonitor() );
				// If destroy not called from self onRunTask()
				if ( ! it->second.isRunning() ) {
					TaskMonitor item = it->second;
					m_mTaskMonitors.erase(it);
					oTaskMonitorLock.unlock();
					item.destroy();
				}
				else
					it->second.setDestroyOnFinish ( true );
			}
		}

		void ThreadPool::schedule(Handler const &pfHandler) {
			boost::recursive_mutex::scoped_lock oLock (m_mxSystemTaskId);
			schedule((*m_itCurrentTaskId), pfHandler);
			if(++m_itCurrentTaskId == m_aSystemTaskIds.end()) {
				m_itCurrentTaskId = m_aSystemTaskIds.begin();
			}
		}

		void ThreadPool::scheduleN(unsigned int nThreadCnt, Handler const & pfHandler) {
			for ( unsigned int i = 0; i < nThreadCnt; i++) {
				schedule ( pfHandler );
			}
		}

		void ThreadPool::schedule(TaskId nTaskId,Handler const & pfHandler) {
			boost::recursive_mutex::scoped_lock oLock (m_mxTaskMonitors);
			TaskMonitors::iterator it = m_mTaskMonitors.find(nTaskId);
			if(it!=m_mTaskMonitors.end()) {
				it->second.incTaskCount();

				boost::recursive_mutex::scoped_lock oStorageLock (m_mxTaskStorage);
				TasksStorageItem oTaskItem (nTaskId, pfHandler);
				m_aTaskStorage.push_back(oTaskItem);
				m_nTaskStorageSize++;
				checkThreadCounts();
				m_evTaskAvailables.notify_one();
			}
		}

		void ThreadPool::checkThreadCounts() {
			if ( m_nTaskStorageSize > m_nCurrentThreadCount && m_nTaskStorageSize < m_nMaxThreadCount ) {
				m_poThreadGroup->create_thread(
					boost::bind(&ThreadPool::onRunTask, this, createSystemTaskId())
				);
				m_nCurrentThreadCount++;
			}
		}

		void ThreadPool::cancel(TaskId nTaskId) {
			boost::recursive_mutex::scoped_lock oMonitorsLock (m_mxTaskMonitors);

			TaskMonitors::iterator itMonitor = m_mTaskMonitors.find(nTaskId);
			if(itMonitor!=m_mTaskMonitors.end()) {
				boost::recursive_mutex::scoped_lock oStorageLock (m_mxTaskStorage);
				for (TaskStorage::iterator it = m_aTaskStorage.begin();it!=m_aTaskStorage.end();) {
					if(it->equalToTaskId(nTaskId)) {
						m_aTaskStorage.erase(it++);
						itMonitor->second.decTaskCount();
						m_nTaskStorageSize--;
					}
					else
						++it;
				}
			}
		}

		void ThreadPool::cancelAll() {
			boost::recursive_mutex::scoped_lock oStorageLock (m_mxTaskStorage);
			m_aTaskStorage.clear();
			m_nTaskStorageSize=0;
			oStorageLock.unlock();

			boost::recursive_mutex::scoped_lock oMonitorsLock (m_mxTaskMonitors);
			for (TaskMonitors::iterator it = m_mTaskMonitors.begin();it!=m_mTaskMonitors.end();it++) {
				boost::recursive_mutex::scoped_lock oTaskMonitorLock ( it->second.getMonitor() );
				it->second.setTaskCount(0);
			}
		}

		void ThreadPool::wait(TaskId nTaskId) {
			boost::recursive_mutex::scoped_lock oMonitorsLock (m_mxTaskMonitors);

			TaskMonitors::iterator it = m_mTaskMonitors.end();
			unsigned int nTaskCount = 0;
			do {
				it = m_mTaskMonitors.find(nTaskId);
				if(it!=m_mTaskMonitors.end()) {
					boost::recursive_mutex::scoped_lock oTaskMonitorLock ( it->second.getMonitor() );
					nTaskCount = it->second.getTaskCount();
					// If wait not called from self onRunTask()
					if ( !it->second.isRunning()) {
						oMonitorsLock.unlock();
						oTaskMonitorLock.unlock();
						oMonitorsLock.lock();
						it = m_mTaskMonitors.find(nTaskId);
					}
					else {
						// Exiting from wait called from onRunTask()
						it = m_mTaskMonitors.end();
					}
				}
			}
			while(it!=m_mTaskMonitors.end() && nTaskCount>0);
		}

		void ThreadPool::waitAll() {
			boost::recursive_mutex::scoped_lock oMonitorsLock (m_mxTaskMonitors);
			boost::recursive_mutex::scoped_lock oStorageLock (m_mxTaskStorage);

			while(!m_aTaskStorage.empty()) {
				oStorageLock.unlock();
				oMonitorsLock.unlock();
				oMonitorsLock.lock();
				oStorageLock.lock();
			}
		}

		void ThreadPool::onRunTask(unsigned int m_nSystemTaskId) {
			boost::recursive_mutex::scoped_lock oStorageLock (m_mxTaskStorage);
			do {
				if ( !m_bStopFlag && m_aTaskStorage.empty()) {
					m_evTaskAvailables.wait(oStorageLock);
				}
				else
				if(!m_bStopFlag) {
					TasksStorageItem oTaskItem = m_aTaskStorage.front();
					m_aTaskStorage.pop_front();
					m_nTaskStorageSize--;
					oStorageLock.unlock();

					boost::recursive_mutex::scoped_lock oMonitorsLock (m_mxTaskMonitors);

					TaskMonitors::iterator it = m_mTaskMonitors.find(oTaskItem.getTaskId());
					if(it!=m_mTaskMonitors.end()) {
						boost::recursive_mutex::scoped_lock oTaskMonitorLock ( it->second.getMonitor() );
						oMonitorsLock.unlock();

						it->second.setRunning ( true );
						oTaskItem.executeHandler();
						it->second.decTaskCount();
						it->second.setRunning ( false );
						bool bDestroyOnFinish = it->second.isDestroyOnFinish();
						oTaskMonitorLock.unlock();
						if ( bDestroyOnFinish ) {
							destroyTaskId(oTaskItem.getTaskId());
						}
						oStorageLock.lock();
					}
				}
			}
			while(!m_bStopFlag);
		}

		const std::string& ThreadPool::getThreadPoolName() {
			return m_sThreadPoolName;
		}
	}
}

